import sys
from colorama import Fore, Style
from enum import Enum
from datetime import datetime
from prompt_toolkit import PromptSession
from prompt_toolkit.styles import Style as PromptStyle
from prompt_toolkit.formatted_text import FormattedText

class OutputType(Enum):
    SYSTEM = "system"      # 系统消息
    CODE = "code"          # 代码
    RESULT = "result"      # 结果
    ERROR = "error"        # 错误
    INFO = "info"          # 提示信息
    PLANNING = "planning"  # 计划
    PROGRESS = "progress"  # 进度
    SUCCESS = "success"    # 成功
    WARNING = "warning"    # 警告
    DEBUG = "debug"        # 调试
    USER = "user"          # 用户输入
    TOOL = "tool"          # 工具输出

class EnhancedOutput:
    """增强版美化输出类"""

    # 新颜色方案
    THEMES = {
        OutputType.SYSTEM: "ansiBrightCyan",      # 亮青色 - AI助手
        OutputType.CODE: "ansiBrightGreen",       # 亮绿色 - 代码
        OutputType.RESULT: "ansiBrightBlue",      # 亮蓝色 - 结果
        OutputType.ERROR: "ansiBrightRed",        # 亮红色 - 错误
        OutputType.INFO: "ansiBrightYellow",      # 亮黄色 - 提示
        OutputType.PLANNING: "ansiBrightMagenta",  # 亮洋红色 - 计划
        OutputType.PROGRESS: "ansiWhite",         # 白色 - 进度
        OutputType.SUCCESS: "ansiGreen",          # 绿色 - 成功
        OutputType.WARNING: "ansiYellow",         # 黄色 - 警告
        OutputType.DEBUG: "ansiBlue",             # 蓝色 - 调试
        OutputType.USER: "ansiGreen",             # 绿色 - 用户
        OutputType.TOOL: "ansiYellow",            # 黄色 - 工具
    }

    # 新图标方案
    DECORATIONS = {
        OutputType.SYSTEM: "⚙️",    # 齿轮 - AI助手
        OutputType.CODE: "💻",      # 电脑 - 代码
        OutputType.RESULT: "🌟",    # 星星 - 结果
        OutputType.ERROR: "💥",     # 爆炸 - 错误
        OutputType.INFO: "📢",      # 喇叭 - 提示
        OutputType.PLANNING: "📅",  # 日历 - 计划
        OutputType.PROGRESS: "⏱️",  # 秒表 - 进度
        OutputType.SUCCESS: "🎯",   # 靶心 - 成功
        OutputType.WARNING: "🚧",   # 施工 - 警告
        OutputType.DEBUG: "🔎",     # 放大镜 - 调试
        OutputType.USER: "👥",      # 人群 - 用户
        OutputType.TOOL: "🧰",      # 工具箱 - 工具
    }

    # 新前缀方案
    TITLES = {
        OutputType.SYSTEM: "Sys",
        OutputType.CODE: "Dev",
        OutputType.RESULT: "Res",
        OutputType.ERROR: "Err",
        OutputType.INFO: "Note",
        OutputType.PLANNING: "Plan",
        OutputType.PROGRESS: "Progress",
        OutputType.SUCCESS: "Done",
        OutputType.WARNING: "Alert",
        OutputType.DEBUG: "Debug",
        OutputType.USER: "User",
        OutputType.TOOL: "Tool",
    }

    @staticmethod
    def decorate(text: str, output_type: OutputType, show_time: bool = True) -> str:
        """装饰输出文本"""
        theme = EnhancedOutput.THEMES.get(output_type, "")
        decoration = EnhancedOutput.DECORATIONS.get(output_type, "")
        title = EnhancedOutput.TITLES.get(output_type, "")

        # 添加时间戳
        time_str = f"[{Fore.WHITE}{datetime.now().strftime('%H:%M:%S')}{Style.RESET_ALL}] " if show_time else ""

        # 格式化输出
        decorated_text = f"{time_str}{Fore.__dict__.get(theme.upper(), Fore.WHITE)}{decoration}  {text}{Style.RESET_ALL}"
        return decorated_text

    @staticmethod
    def display(text: str, output_type: OutputType, show_time: bool = True):
        """显示装饰后的文本"""
        print(EnhancedOutput.decorate(text, output_type, show_time))
        if output_type == OutputType.ERROR:
            import traceback
            EnhancedOutput.display(f"Error trace: {traceback.format_exc()}", OutputType.INFO)

    @staticmethod
    def showHeader(title: str, output_type: OutputType = OutputType.INFO):
        """显示分段标题"""
        width = 100
        theme = EnhancedOutput.THEMES.get(output_type, "")
        print(f"\n{Fore.__dict__.get(theme.upper(), Fore.WHITE)}" + "=" * width + f"{Style.RESET_ALL}")
        EnhancedOutput.display(title.center(width - 25), output_type, show_time=False)
        print(f"{Fore.__dict__.get(theme.upper(), Fore.WHITE)}" + "=" * width + f"{Style.RESET_ALL}\n")

    @staticmethod
    def stream(text: str):
        """流式输出，不换行"""
        theme = EnhancedOutput.THEMES.get(OutputType.SYSTEM, "")
        sys.stdout.write(f"{Fore.__dict__.get(theme.upper(), Fore.WHITE)}{text}{Style.RESET_ALL}")
        sys.stdout.flush()

    @staticmethod
    def endStream():
        """结束流式输出，换行"""
        sys.stdout.write("\n")
        sys.stdout.flush()

def inputLine(tip: str = "") -> str:
    sessions= PromptSession(history=None)
    style = PromptStyle.from_dict({
        "prompt": "ansiBrightGreen"
    })
    return sessions.prompt(f"{tip}", style=style)

def inputMultiline(tip:str = "") -> str:
    print(f"{Fore.GREEN}{tip}{Style.RESET_ALL}")
    session = PromptSession(history=None)
    style = PromptStyle.from_dict({
        "prompt": "ansiBrightGreen"
    })
    lines=[]
    try:
        while True:
            prompt = FormattedText([
                ('class:prompt', '... ' if lines else '>>> ')
            ])
            line = session.prompt(prompt, style=style).trip()

            if not line:
                if not lines:
                    return                
            lines.append(line)
    except KeyboardInterrupt:
        EnhancedOutput.display("\n Input canceled", OutputType.INFO, show_time=False)
        return "__interrupt__"
    except EOFError:
        EnhancedOutput.display("\n Input canceled", OutputType.INFO, show_time=False)   
        return "__interrupt__"
    

    def getMaxContextLenght():
        return '131072' # 默认128k